#include <algorithm>
#include <assert.h>
#include <chrono>
#include <climits>
#include <iostream>
#include <omp.h>
#include <stdlib.h>
#include <time.h>
#include <vector>

using std::cout;
using std::endl;
using std::vector;
using std::chrono::duration_cast;
using std::chrono::high_resolution_clock;
using std::chrono::nanoseconds;

#define DEBUG
// #undef DEBUG

bool isAscend(vector<int>& nums) {
	int n = nums.size();
	for (int i = 0; i < n - 1; i++)
		if (nums[i] > nums[i + 1])
			return false;
	return true;
}

void swap(int& left, int& right) {
	int tmp = left;
	left = right;
	right = tmp;
}

int quickPartition(vector<int>& nums, int lo, int hi) { // [lo, hi]
	int p = nums[hi];
	int i = lo, j = lo;
	while (j < hi) {
		if (nums[j] <= p) {
			swap(nums[i], nums[j]);
			i += 1;
		}
		j += 1;
	}
	swap(nums[i], nums[hi]);
	return i;
}

void quickSort(vector<int>& nums, int lo, int hi) { // [lo, hi]
	if (lo >= hi) return;
	int p = quickPartition(nums, lo, hi);
	quickSort(nums, lo, p - 1);
	quickSort(nums, p + 1, hi);
}

int main(int argc, char* argv[]) {
	// get info
    if (argc != 3) {
        cout << "Usage: " << argv[0] << " arraysize numprocs" << endl;
        return 1; 
    }
	int n = atoi(argv[1]);
	int numprocs = atoi(argv[2]);
    assert(n > 0 && numprocs > 0);

	// init nums
	vector<int> nums(n);
	srand((unsigned)time(NULL));
	for (int i = 0; i < n; i++) nums[i] = rand();

    // quick sort
    auto starttime = high_resolution_clock::now();
    if (numprocs == 1) {
        // std::sort(nums.begin(), nums.end());
        quickSort(nums, 0, n - 1);
    } else {
        // partition
        vector<vector<int>> omp_nums(numprocs);
        int n_per = n / numprocs, n_last = n - n_per * numprocs;
        for (int i = 0; i < numprocs; i++) {
            if (i < numprocs - 1) {
                omp_nums[i].assign(nums.begin() + i * n_per, nums.begin() + i * n_per + n_per);
            } else {
                omp_nums[i].assign(nums.begin() + i * n_per, nums.begin() + i * n_per + n_last);
            }
        }
        // parallel quick sort
	    omp_set_num_threads(numprocs);
        #pragma omp parallel for
	        for (int i = 0; i < numprocs; i++)
                // std::sort(omp_nums[i].begin(), omp_nums[i].end());
                quickSort(omp_nums[i], 0, omp_nums[i].size() - 1);
        // merge
        vector<int> ompn(numprocs, n_per);
        ompn[numprocs - 1] = n_last;
        vector<int> ptr(numprocs, 0);
        bool flag = true;
        int j = 0, minv, mini;
        while (j < n) {
            minv = INT_MAX;
            for (int i = 0; i < numprocs; i++) {
                if (ptr[i] < ompn[i] && omp_nums[i][ptr[i]] < minv) {
                    minv = omp_nums[i][ptr[i]];
                    mini = i;
                }
            }
            nums[j++] = minv;
            ptr[mini] += 1;
        }
    }
    auto stoptime = high_resolution_clock::now();
    auto elapsedS = duration_cast<nanoseconds>(stoptime - starttime).count() / 1e+9;
    cout << "arraysize = " << n << ", numprocs = " << numprocs << ", time = " << elapsedS << " s" << endl;

#ifdef DEBUG
    assert(isAscend(nums));
#endif

	return 0;
}
